package org.springblade.metadata.service.impl;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.apache.atlas.AtlasClientV2;
import org.apache.atlas.AtlasServiceException;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
import org.apache.commons.collections.CollectionUtils;
import org.apache.curator.shaded.com.google.common.collect.Sets;
import org.apache.atlas.model.discovery.AtlasSearchResult;
import org.apache.atlas.model.discovery.SearchParameters;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasEntityHeader;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.springblade.metadata.service.AtlasProxyService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.raysdata.atlas.common.SystemConstant;

@Component
public class EntityService {

	private final AtlasProxyService atlasProxyService;

	private final AtlasClientV2 client;

	@Autowired
	public EntityService(AtlasProxyService atlasProxyService) {
		super();
		this.atlasProxyService = atlasProxyService;
		this.client = this.atlasProxyService.get();
	}

	public AtlasEntity save(AtlasEntityWithExtInfo entity) throws AtlasServiceException {

		AtlasEntity ret = null;
		EntityMutationResponse response = client.createEntity(entity);
		List<AtlasEntityHeader> entities = response.getEntitiesByOperation(EntityOperation.CREATE);
		if (CollectionUtils.isEmpty(entities)) {
			entities = response.getEntitiesByOperation(EntityOperation.UPDATE);
		}
		if (CollectionUtils.isNotEmpty(entities)) {
			AtlasEntityWithExtInfo getByGuidResponse = client.getEntityByGuid(entities.get(0).getGuid());

			ret = getByGuidResponse.getEntity();

			System.out.println("Created entity of type [" + ret.getTypeName() + "], guid: " + ret.getGuid());
		}

		return ret;
	}

	public AtlasEntity get(String guid) throws AtlasServiceException {
		AtlasEntity ret = null;

		AtlasEntityWithExtInfo getByGuidResponse = client.getEntityByGuid(guid);

		ret = getByGuidResponse.getEntity();

		System.out.println("Created entity of type [" + ret.getTypeName() + "], guid: " + ret.getGuid());

		return ret;
	}

	public void addLable(String guid, List<String> lables) throws AtlasServiceException {

		Set<String> ls = new HashSet<>(lables);
		client.addLabels(guid, ls);
	}

	public void addLable(String guid, Set<String> lables) throws AtlasServiceException {
		client.addLabels(guid, lables);
	}

	public void setLable(String guid, List<String> lables) throws AtlasServiceException {

		Set<String> ls = new HashSet<>(lables);
		client.setLabels(guid, ls);
	}

	public void setLable(String guid, Set<String> lables) throws AtlasServiceException {
		client.addLabels(guid, lables);
	}

	public AtlasSearchResult search(SearchParameters searchParameters) throws AtlasServiceException {

		searchParameters.setTermName(searchParameters.getTermName() + "@" + SystemConstant.GLOSSARY_GTOUP_NAME);

		searchParameters.setExcludeDeletedEntities(true);

		searchParameters.setIncludeClassificationAttributes(true);

		searchParameters.setIncludeSubClassifications(true);

		return client.facetedSearch(searchParameters);
	}

	public AtlasEntityHeader delete(String guid) throws AtlasServiceException {
		EntityMutationResponse rep = client.deleteEntityByGuid(guid);

		List<AtlasEntityHeader> entities = rep.getEntitiesByOperation(EntityOperation.DELETE);
		AtlasEntityHeader eh = entities.get(0);
		return eh;
	}
	
	public void pure(String guid) throws AtlasServiceException {
		client.purgeEntitiesByGuids(Sets.newHashSet(guid));
	}
}
